import ast
import consts
from exception import TypeJsonParseError


class TypeTag:  # (str, TypeSymbol or GenericTypeSymbol)
    def __init__(self, name, the_ast):
        self.name = name
        self.ast = the_ast

    def __str__(self):
        return self.name


class Scope:
    def __init__(self, parents):
        tmp = {}
        res = []
        for x in parents:
            if x not in tmp:
                res.append(x)
                tmp[x] = 1
        self.parents = res
        self.types = {}

    def __get_type_with_name_not_raise(self, name, line_col):
        if name in self.types:
            return self.types[name]
        for parent in self.parents:
            res = parent.__get_type_with_name_not_raise(name, line_col)
            if res:
                return res
        return None

    def get_type_with_name(self, name, line_col):
        res = self.__get_type_with_name_not_raise(name, line_col)
        if res:
            return res
        raise TypeJsonParseError('formatter', 'unknown type tag ' + str(name), line_col)

    def add_type_tag(self, tag, line_col):
        if tag.name in self.types:
            raise TypeJsonParseError('formatter', 'duplicated type tag ' + tag.name, line_col)
        self.types[tag.name] = tag


class Formatter:
    def __init__(self):
        # init scope and built in tags
        scope = self.scope = Scope([])

        # built-in type symbols
        type_symbol_any = ast.TypeSymbol('Any', None, None)
        type_symbol_any.built_in = True
        type_symbol_number = ast.TypeSymbol('Number', None, None)
        type_symbol_number.built_in = True
        type_symbol_int = ast.TypeSymbol('Int', None, None)
        type_symbol_int.built_in = True
        type_symbol_float = ast.TypeSymbol('Float', None, None)
        type_symbol_float.built_in = True
        type_symbol_string = ast.TypeSymbol('String', None, None)
        type_symbol_string.built_in = True
        type_symbol_bool = ast.TypeSymbol('Bool', None, None)
        type_symbol_bool.built_in = True
        type_symbol_list = ast.TypeSymbol('List',
                                          [
                                              ast.GenericTypeSymbol('E', None, None)
                                          ], None)
        type_symbol_list.built_in = True
        type_symbol_set = ast.TypeSymbol('Set',
                                         [
                                             ast.GenericTypeSymbol('E', None, None)
                                         ], None)
        type_symbol_set.built_in = True

        # built-in tags
        tag_any = TypeTag('Any', type_symbol_any)
        tag_number = TypeTag('Number', type_symbol_number)
        tag_int = TypeTag('Int', type_symbol_int)
        tag_float = TypeTag('Float', type_symbol_float)
        tag_string = TypeTag('String', type_symbol_string)
        tag_bool = TypeTag('Bool', type_symbol_bool)
        tag_list = TypeTag('List', type_symbol_list)
        tag_set = TypeTag('Set', type_symbol_set)

        # assign tags to ast
        type_symbol_any.type_tag = tag_any
        type_symbol_number.type_tag = tag_number
        type_symbol_int.type_tag = tag_int
        type_symbol_float.type_tag = tag_float
        type_symbol_string.type_tag = tag_string
        type_symbol_bool.type_tag = tag_bool
        type_symbol_list.type_tag = tag_list
        type_symbol_set.type_tag = tag_set

        # assign tags to generics
        tag_list_e = TypeTag('E', type_symbol_list.generic_type_symbols[0])
        type_symbol_list.generic_type_symbols[0].type_tag = tag_list_e

        tag_set_e = TypeTag('E', type_symbol_set.generic_type_symbols[0])
        type_symbol_set.generic_type_symbols[0].type_tag = tag_set_e

        # generate def
        def_any = ast.Def(type_symbol_any, None, [], None)
        type_ref_symbol_any = ast.TypeRefSymbol('Any', None, None)
        def_number = ast.Def(type_symbol_number, type_ref_symbol_any, [], None)
        type_ref_symbol_number = ast.TypeRefSymbol('Number', None, None)
        def_int = ast.Def(type_symbol_int, type_ref_symbol_number, [], None)
        def_float = ast.Def(type_symbol_float, type_ref_symbol_number, [], None)
        def_string = ast.Def(type_symbol_string, type_ref_symbol_any, [], None)
        def_bool = ast.Def(type_symbol_bool, type_ref_symbol_any, [], None)
        def_list = ast.Def(type_symbol_list, type_ref_symbol_any, [], None)
        def_set = ast.Def(type_symbol_set, type_ref_symbol_any, [], None)

        # fill parent_type_tag | a_def
        def_any.parent_type_tag = None
        def_number.parent_type_tag = tag_any
        def_int.parent_type_tag = tag_number
        def_float.parent_type_tag = tag_number
        def_string.parent_type_tag = tag_any
        def_bool.parent_type_tag = tag_any
        def_list.parent_type_tag = tag_any
        def_set.parent_type_tag = tag_set

        type_symbol_any.a_def = def_any
        type_symbol_number.a_def = def_number
        type_symbol_int.a_def = def_int
        type_symbol_float.a_def = def_float
        type_symbol_string.a_def = def_string
        type_symbol_bool.a_def = def_bool
        type_symbol_list.a_def = def_list
        type_symbol_set.a_def = def_set

        # register into scope
        self.tag_any = tag_any
        scope.types['Any'] = tag_any
        scope.types['Number'] = tag_number
        scope.types['Int'] = tag_int
        scope.types['Float'] = tag_float
        scope.types['String'] = tag_string
        scope.types['Bool'] = tag_bool
        scope.types['List'] = tag_list
        scope.types['Set'] = tag_set

    def format(self, ast_list, old_ast):
        name_ast_map = {}
        # extract all defined type tags
        count = 0
        for a in ast_list:
            module = a[consts.AST_MODULE]
            if module:
                module_name = module.name.string
            else:
                module_name = '$the$generated$module$name' + str(count)
                count += 1
            if module_name in name_ast_map:
                if module:
                    raise TypeJsonParseError('formatter', 'duplicated module name ' + module_name, module.line_col)
                else:
                    raise Exception('the generated module name is already used! ' + module_name)
            name_ast_map[module_name] = a
            scope = Scope([self.scope])
            a['scope'] = scope
            for d in a[consts.AST_DEF]:
                type_symbol = d.type_symbol
                type_symbol.a_def = d
                type_name = type_symbol.type_name
                tag = TypeTag(type_name, type_symbol)
                # add into scope
                scope.add_type_tag(tag, type_symbol.line_col)
                # add into ast
                type_symbol.type_tag = tag
        # for each module ast, import the `use` modules
        for a in ast_list:
            parent_scopes = []
            use_list = a[consts.AST_USE]
            for use in use_list:
                name = use.name.string
                if name in name_ast_map:
                    aa = name_ast_map[name]
                    scope = aa['scope']
                    parent_scopes.append(scope)
            parent_scopes.append(a['scope'])
            a['final_scope'] = Scope(parent_scopes)
        # for each ast structure, add type_tag to it
        for a in ast_list:
            scope = a['final_scope']
            for d in a[consts.AST_DEF]:
                # def scope
                d_scope = Scope([scope])
                type_symbol = d.type_symbol
                # type tags on type symbol
                self.__tag_type_symbol(type_symbol, d_scope)
                # type tags on parent
                parent_type_ref_symbol = d.parent_type_ref_symbol
                if parent_type_ref_symbol:
                    self.__tag_real_type(parent_type_ref_symbol, d_scope)
                    # the parent type (not type params) should point to a def
                    if isinstance(parent_type_ref_symbol.type_tag.ast, ast.GenericTypeSymbol):
                        raise TypeJsonParseError('formatter', 'type parameter is not allowed for parent type',
                                                 parent_type_ref_symbol.line_col)
                # type tags on members
                for member in d.def_members:
                    s = Scope([d_scope])
                    self.__tag_real_type(member.optional_type_symbol.type_ref_symbol, s)
            o = a[consts.AST_VALUE]
            if o:
                self.__tag_value(o, scope)
        # form parents
        for a in ast_list:
            for d in a[consts.AST_DEF]:
                if d.parent_type_ref_symbol:
                    d.parent_type_tag = d.parent_type_ref_symbol.type_tag
                else:
                    d.parent_type_tag = self.tag_any
        return ast_list + old_ast

    def __tag_value(self, type_value, scope):
        type_ref_symbol = type_value.type_ref_symbol
        if type_ref_symbol:
            self.__tag_real_type(type_ref_symbol, scope)
        value = type_value.value
        if isinstance(value, ast.Object):
            for member in value.members:
                self.__tag_value(member.value, scope)
        elif isinstance(value, ast.Array):
            for elem in value.elements:
                self.__tag_value(elem, scope)
        elif isinstance(value, ast.Set):
            for elem in value.elements:
                self.__tag_value(elem, scope)

    def __tag_type_symbol(self, type_symbol, scope):
        # process type tags
        type_name = type_symbol.type_name
        tag = scope.get_type_with_name(type_name, type_symbol.line_col)
        type_symbol.type_tag = tag
        if type_symbol.generic_type_symbols:
            tags = []
            for g in type_symbol.generic_type_symbols:
                tag = TypeTag(g.type_name, g)
                g.type_tag = tag
                t_scope = Scope([scope])
                t_scope.add_type_tag(tag, g.line_col)
                tags.append(tag)
                if g.type_ref_symbol:
                    self.__tag_real_type(g.type_ref_symbol, t_scope)
                    g.type_symbol_tag = g.type_ref_symbol.type_tag
                else:
                    g.type_symbol_tag = self.tag_any
            for tag in tags:
                scope.add_type_tag(tag, type_symbol.line_col)

    @staticmethod
    def __tag_real_type(type_ref_symbol, scope):
        type_name = type_ref_symbol.type_name
        tag = scope.get_type_with_name(type_name, type_ref_symbol.line_col)
        type_ref_symbol.type_tag = tag
        if type_ref_symbol.type_ref_symbols:
            for type_ref_symbol in type_ref_symbol.type_ref_symbols:
                Formatter.__tag_real_type(type_ref_symbol, scope)
